"""
检索elasticsearch 告警
"""


class OP:
    """数据对比"""
    @staticmethod
    def lt(v1, v2):
        return v1 < v2

    @staticmethod
    def lte(v1, v2):
        return v1 <= v2

    @staticmethod
    def eq(v1, v2):
        return v1 == v2

    @staticmethod
    def gt(v1, v2):
        return v1 > v2

    @staticmethod
    def gte(v1, v2):
        return v1 >= v2


class Alert:
    """
    查询数据
    检查数据告警
    触发告警
    """

    def __init__(self, client, q, triggers, sample_size=10):
        self.client = 'session'
        self.q = q

        self.triggers = triggers
        self.sample_size = sample_size
        self.resp = None

    def build_query(self):
        pass

    def query(self):
        self.build_query()
        resp = self.client.execute()
        self.resp = resp

    def check_trigger(self):
        self.query()
        for trigger in self.triggers:
            v = self.resp.hits.total.value
            if op_func := getattr(OP, trigger['op']):
                if op_func(v, trigger['value']):
                    return v, trigger
        return False


class CountAlert(Alert):
    """分组"""
    def __init__(self, client, q, triggers, sample_size=10):
        super().__init__(client, q, triggers, sample_size=sample_size)

    def check_trigger(self):
        self.query()
        for trigger in self.triggers:
            v = self.resp.hits.total.value
            if op_func := getattr(OP, trigger['op']):
                if op_func(v, trigger['value']):
                    return v, trigger
        return False


class GroupAlert(Alert):
    """分组告警"""
    def __init__(self, client, q, groups, triggers, sample_size=10):
        self.groups = groups
        super().__init__(client, q, triggers, sample_size=sample_size)

    def build_query(self):
        sources = []
        for group in self.groups:
            agg_type, field = group['name_or_agg'], group['field']
            sources.append({f'{agg_type}_{field}': A(**group)})
        composite_agg = A('composite', sources=sources)
        for trigger in self.triggers:
            agg_type, field = trigger['metric']['name_or_agg'], trigger['metric']['field']
            composite_agg.bucket(f'{agg_type}_{field}', A(**trigger['metric']))

        self.client.aggs.bucket('composite', composite_agg)

    def check_trigger(self):
        self.query()

        group_metrics = defaultdict()
        for bucket in self.resp.aggregations.composite.buckets:
            bucket = bucket.to_dict()
            result = {k: v['value'] for k, v in bucket.items() if isinstance(v, dict) and 'value' in v}
            result['doc_count'] = bucket['doc_count']
            group_metrics[tuple(bucket['key'].values())] = result

        for group, metric in group_metrics.items():
            for trigger in self.triggers:
                if op_func := getattr(OP, trigger['op']):
                    agg_type, field = trigger['metric']['name_or_agg'], trigger['metric']['field']
                    v = metric[f'{agg_type}_{field}']
                    if op_func(v, trigger['value']):
                        yield group, v, trigger


class MetricAlert(Alert):

    def build_query(self):
        for trigger in self.triggers:
            agg_type, field = trigger['metric']['name_or_agg'], trigger['metric']['field']
            self.client.aggs.bucket(f'{agg_type}_{field}', A(**trigger['metric']))

    def check_trigger(self):
        self.query()
        for trigger in self.triggers:
            agg_type, field = trigger['metric']['name_or_agg'], trigger['metric']['field']
            v = getattr(self.resp.aggregations, f'{agg_type}_{field}').value
            if op_func := getattr(OP, trigger['op']):
                if op_func(v, trigger['value']):
                    return v, trigger
        return False


if __name__ == '__main__':
    from common.package.elastic_db import es_db
    es_db.init(hosts='http://192.168.101.79:9200')

    index = 'dsm-1086afd495bd11efae39622bfe3b1ffc'
    raw = {'query': {'match_all': {}}, 'size': 10}

    def count_test():
        triggers = [
            {'metric': {'name_or_agg': 'hits', 'field': 'total'}, 'op': 'gt', 'value': 1}
        ]
        count_alert = CountAlert(es_db.client, raw, triggers)
        result = count_alert.check_trigger()
        if not result:
            return False
        print(result)


    def metric_test():
        triggers = [
            {'metric': {'name_or_agg': 'cardinality', 'field': '_host'}, 'op': 'gt', 'value': 1, 'level': 3},
            {'metric': {'name_or_agg': 'value_count', 'field': '_host'}, 'op': 'gt', 'value': 1, 'level': 2},
        ]
        count_alert = MetricAlert(es_db.client, raw, triggers)
        result = count_alert.check_trigger()
        if not result:
            return False
        print(result)

    def group_test():
        groups = [
            {'name_or_agg': 'terms', 'field': '_sourceid'},
            {'name_or_agg': 'terms', 'field': '_host'}]
        triggers = [
            {'metric': {'name_or_agg': 'cardinality', 'field': '_host'}, 'op': 'gt', 'value': 1, 'level': 3},
            {'metric': {'name_or_agg': 'value_count', 'field': '_host'}, 'op': 'gt', 'value': 1, 'level': 2},
        ]
        count_alert = GroupAlert(es_db.client, raw, groups, triggers)
        result = list(count_alert.check_trigger())
        if not result:
            return False
        print(result)

    # count_test()
    # metric_test()
    group_test()
